import re
import matplotlib.pyplot as plt

# Step 1: 读取 log 文件
log_file = "log.log"  # 替换为你的 log 文件名

with open(log_file, "r") as file:
    log_data = file.readlines()

# Step 2: 提取 iteration、actor/pg_loss 和 response_length/mean 数据
iterations = []
pg_losses = []
response_lengths_mean = []

for line in log_data:
    # 使用正则表达式提取 iteration、actor/pg_loss 和 response_length/mean
    iteration_match = re.search(r"iteration: (\d+)", line)
    pg_loss_match = re.search(r"actor/pg_loss\s*:\s*([-+]?\d*\.\d+|\d+)", line)
    response_length_mean_match = re.search(r"response_length/mean\s*:\s*([-+]?\d*\.\d+|\d+)", line)

    if iteration_match:
        iterations.append(int(iteration_match.group(1)))  # 提取 iteration
    if pg_loss_match:
        pg_losses.append(float(pg_loss_match.group(1)))  # 提取 actor/pg_loss
    if response_length_mean_match:
        response_lengths_mean.append(float(response_length_mean_match.group(1)))  # 提取 response_length/mean

# Step 3: 绘制 actor/pg_loss 折线图
plt.figure(figsize=(10, 6))
plt.plot(iterations[:len(pg_losses)], pg_losses, marker="o", linestyle="-", color="b", label="actor/pg_loss")
plt.title("Actor/PG Loss Over Iterations", fontsize=16)
plt.xlabel("Iteration", fontsize=14)
plt.ylabel("Actor/PG Loss", fontsize=14)
plt.grid(True, linestyle="--", alpha=0.6)
plt.legend(fontsize=12)
plt.tight_layout()
# 保存 actor/pg_loss 图表
pg_loss_output_file = "actor_pg_loss_plot.png"
plt.savefig(pg_loss_output_file, dpi=300)
print(f"Actor/PG Loss 图表已保存为 {pg_loss_output_file}")
plt.close()

# Step 4: 绘制 response_length/mean 折线图
plt.figure(figsize=(10, 6))
plt.plot(iterations[:len(response_lengths_mean)], response_lengths_mean, marker="o", linestyle="-", color="g",
         label="response_length/mean")
plt.title("Response Length Mean Over Iterations", fontsize=16)
plt.xlabel("Iteration", fontsize=14)
plt.ylabel("Response Length Mean", fontsize=14)
plt.grid(True, linestyle="--", alpha=0.6)
plt.legend(fontsize=12)
plt.tight_layout()
# 保存 response_length/mean 图表
response_length_output_file = "response_length_mean_plot.png"
plt.savefig(response_length_output_file, dpi=300)
print(f"Response Length Mean 图表已保存为 {response_length_output_file}")
plt.close()